package com.liru.neo4j_demo.service;

import com.liru.neo4j_demo.model.Entity;
import com.liru.neo4j_demo.model.Relation;
import dev.langchain4j.internal.Json;
import dev.langchain4j.model.openai.OpenAiChatModel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 自然语言处理服务类，使用语言模型替换正则表达式进行实体和关系抽取
 */
@Service
public class NLPProcessingService {

    // 日志记录器
    private static final Logger log = LoggerFactory.getLogger(NLPProcessingService.class);

    // OpenAI聊天模型
    private final OpenAiChatModel openAiChatModel;

    /**
     * 构造函数，注入依赖
     * @param openAiChatModel OpenAI聊天模型
     */
    @Autowired
    public NLPProcessingService(OpenAiChatModel openAiChatModel) {
        this.openAiChatModel = openAiChatModel;
    }

    /**
     * 使用语言模型从文本中抽取实体
     * @param text 待处理的文本
     * @param knowledgeId 知识ID
     * @return 实体列表
     */
    public List<Entity> extractEntities(String text, String knowledgeId) {
        List<Entity> entities = new ArrayList<>();

        // 构建提示词
        String prompt = String.format(
            "从以下文本中提取实体，并返回JSON格式：\n" +
            "{\n" +
            "  \"entities\": [\n" +
            "    {\"name\": \"实体名称\", \"type\": \"类型\", \"properties\": {}}\n" +
            "  ]\n" +
            "}\n\n" +
            "文本：%s",
            text
        );

        // 调用语言模型
        String response = openAiChatModel.chat(prompt);
        System.out.println("模型响应: " + response);  // 调试输出
        log.info("实体抽取模型响应: {}", response);

        // 解析模型响应
        try {
            // 使用JSON库解析响应
            Object parsedResponse = Json.fromJson(response.substring(8, response.length()-4), Object.class);
            if (parsedResponse instanceof Map) {
                Map<String, Object> jsonResponse = (Map<String, Object>) parsedResponse;
                Object entitiesObj = jsonResponse.get("entities");
                if (entitiesObj instanceof List) {
                    List<?> rawEntityList = (List<?>) entitiesObj;
                    Map<String, Entity> entityMap = new HashMap<>();

                    for (Object item : rawEntityList) {
                        if (item instanceof Map) {
                            Map<String, Object> entityData = (Map<String, Object>) item;
                            Object nameObj = entityData.get("name");
                            Object typeObj = entityData.get("type");
                            Object propertiesObj = entityData.get("properties");

                            if (nameObj instanceof String && typeObj instanceof String) {
                                String name = (String) nameObj;
                                String type = (String) typeObj;
                                Map<String, Object> properties = new HashMap<>();

                                if (propertiesObj instanceof Map) {
                                    properties = (Map<String, Object>) propertiesObj;
                                }

                                if (!entityMap.containsKey(name)) {
                                    Entity entity = new Entity();
                                    entity.setName(name);
                                    entity.setType(type);
                                    entity.setKnowledgeId(knowledgeId);
                                    entity.setProperties(properties);
                                    entityMap.put(name, entity);
                                }
                            }
                        }
                    }
                    entities.addAll(entityMap.values());
                    log.info("实体抽取结果：{}", entities);
                } else {
                    System.err.println("'entities' 字段不是列表类型");
                    log.error("实体抽取错误: 'entities' 字段不是列表类型");
                }
            } else {
                System.err.println("响应不是有效的JSON对象");
                log.error("实体抽取错误: 响应不是有效的JSON对象");
            }
        } catch (Exception e) {
            // 处理解析错误
            System.err.println("实体抽取失败: " + e.getMessage());
            log.error("实体抽取失败: {}", e.getMessage(), e);
        }

        return entities;
    }

    /**
     * 使用语言模型从文本中抽取实体关系
     * @param text 待处理的文本
     * @param entities 实体列表
     * @param knowledgeId 知识ID
     * @return 关系列表
     */
    public List<Relation> extractRelations(String text, List<Entity> entities, String knowledgeId) {
        List<Relation> relations = new ArrayList<>();

        // 构建实体列表字符串
        StringBuilder entitiesStr = new StringBuilder();
        for (Entity entity : entities) {
            entitiesStr.append("- ").append(entity.getName()).append(" (").append(entity.getType()).append(")\n");
        }

        // 构建提示词
        String prompt = "从以下文本中提取实体之间的关系，仅考虑提供的实体列表中的实体，并返回JSON格式：\n" +
                "{\"relations\": [\n" +
                "  {\"source\": \"源实体名称\", \"target\": \"目标实体名称\", \"type\": \"关系类型\", \"properties\": {\"属性名\": \"属性值\"}}\n" +
                "]}\n\n" +
                "实体列表：\n" + entitiesStr + "\n\n" +
                "文本：" + text;

        // 调用语言模型
        String response = openAiChatModel.chat(prompt);
        log.info("关系抽取模型响应: {}", response);

        // 解析模型响应
        try {
            // 使用JSON库解析响应
            Map<String, Object> jsonResponse = Json.fromJson(response.substring(8, response.length()-4), Map.class);
            List<Map<String, Object>> relationList = (List<Map<String, Object>>) jsonResponse.get("relations");

            if (relationList != null) {
                for (Map<String, Object> relationData : relationList) {
                    String sourceName = (String) relationData.get("source");
                    String targetName = (String) relationData.get("target");
                    String type = (String) relationData.get("type");
                    Map<String, Object> properties = (Map<String, Object>) relationData.getOrDefault("properties", new HashMap<>());

                    // 查找源实体和目标实体
                    Entity source = findEntityByName(entities, sourceName);
                    Entity target = findEntityByName(entities, targetName);

                    if (source != null && target != null) {
                        Relation relation = new Relation();
                        relation.setType(type);
                        relation.setKnowledgeId(knowledgeId);
                        relation.setSource(source);
                        relation.setTarget(target);
                        relation.setProperties(properties);

                        relations.add(relation);
                    }
                }
            }
        } catch (Exception e) {
            // 处理解析错误
            System.err.println("关系抽取失败: " + e.getMessage());
        }
        log.info("关系抽取结果：{}", relations);

        return relations;
    }

    /**
     * 根据名称查找实体
     * @param entities 实体列表
     * @param name 实体名称
     * @return 找到的实体，或null
     */
    private Entity findEntityByName(List<Entity> entities, String name) {
        for (Entity entity : entities) {
            if (entity.getName().equals(name)) {
                return entity;
            }
        }
        return null;
    }
}